#include <iostream>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <string.h>
#include <vector>
#include "opencv2/dnn.hpp"
#include "opencv2/imgproc.hpp"

using namespace cv;
using namespace cv::dnn;
using namespace std;
#include <stdlib.h>
#include <fstream>
#include <cstdlib>
#include <ctime>

void getMaxClass(dnn::Blob &probBlob, int *classId, double *classProb)
{
    Mat probMat = probBlob.matRefConst().reshape(1, 1); //reshape the blob to 1x1000 matrix
    Point classNumber;
    minMaxLoc(probMat, NULL, classProb, NULL, &classNumber);
    *classId = classNumber.x;
}

int recognize(Mat source_image,vector<cv::Mat> images,vector<cv::Rect> rects, string savePath)
{
	String modelTxt = "eagLeNet_deploy.prototxt";
	String modelBin = "eagLeNet_quick_iter_7500.caffemodel";
    //! [Create the importer of Caffe model] 导入一个caffe模型接口 
	Ptr<dnn::Importer> importer; 
	importer = dnn::createCaffeImporter(modelTxt, modelBin);

	if (!importer)
	{
		std::cerr << "Can't load network by using the following files: " << std::endl;
		std::cerr << "prototxt:   " << modelTxt << std::endl;
		std::cerr << "caffemodel: " << modelBin << std::endl;
		exit(-1);
	}

    //! [Initialize network] 通过接口创建和初始化网络
	Net net;
	importer->populateNet(net);  
	importer.release();

	int plane=0;

	int k=0;
	for(int i=0;i<images.size();i++)
	{
		Mat img=images[i];

	    resize(img, img, Size(64, 64));                   //[<Important>]Mnist accepts only 28x28 RGB-images

	    dnn::Blob inputBlob = cv::dnn::Blob(img);   //Convert Mat to dnn::Blob batch of images

	    //! [Set input blob] 将blob输入到网络
	    net.setBlob(".data", inputBlob);        //set the network input

	    //! [Make forward pass] 进行前向传播
	    net.forward();                          //compute output

	    //! [Gather output] 获取概率值
	    dnn::Blob prob = net.getBlob("prob");   //[<Important>] gather output of "prob" layer
	    int classId;
	    double classProb;
	    getMaxClass(prob, &classId, &classProb);//find the best class

	    //! [Print results] 输出结果
	    //std::cout << "Best class: #" << classId << "'" << std::endl;
	    //std::cout << "Probability: " << classProb * 100 << "%" << std::endl;
	    if(classId==0&&classProb*100>90)
	    {
	    	k++;
	    	cout<<classProb*100<<"%"<<" "<<rects[i].tl()<<endl;
	    	cv::rectangle(source_image,rects[i],Scalar(0,0,255),2,1,0);
	    	imwrite("data/img_"+to_string(k)+".jpg",img);
	    	plane++;
	    }
	}
//	imshow("matImage",source_image); 
	imwrite(savePath,source_image);
	waitKey();
	return plane;
}


int main(int argc, char** argv){

/*
*	./a /home/wyj/caffe/examples/eagLeNet/airport1.png 128 128 2
*/
	time_t start=time(NULL);
	cv::Mat image;
	cv::Mat source_image;
	vector<cv::Mat>  images;
	vector<cv::Rect> rects;
	string source=(argc > 1) ? argv[1] : "/home/wyj/caffe/examples/eagLeNet/airport1.png";
	int width=(argc > 2) ? atoi(argv[2]) : 128;
	int height=(argc > 3) ? atoi(argv[3]) : 128;
	int stride=(argc > 4) ? atoi(argv[4]) : 2;
	string savePath=source+"_result.jpg";

	source_image= cv::imread(source);
	cout<<source_image.cols<<"*"<<source_image.rows<<endl;

	for(int i=0 ; i+height< source_image.rows; i+=height/stride)
	{
		for(int j=0 ; j+width<source_image.cols; j+=width/stride)
		{
			image = source_image(cv::Rect(j,i,width,height)).clone();
			images.push_back(image);
			rects.push_back(cv::Rect(j,i,width,height));
		}
	}

	if(images.empty())
	{
		cout<<"empty";
		return 0;
	}
	
	cout<<recognize(source_image,images,rects,savePath)<<endl;
	time_t end=time(NULL);
	cout<<"time: "<<end-start<<" s"<<endl;
	return 0;
}

